package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/shirou/gopsutil/v4/net"
)

// Connection represents a network connection of a process
type Connection struct {
	Protocol   string `json:"protocol"`
	LocalIP    string `json:"local_ip"`
	LocalPort  int    `json:"local_port"`
	RemoteIP   string `json:"remote_ip"`
	RemotePort int    `json:"remote_port"`
	Status     string `json:"status"`
}

// NetStats represents network traffic statistics
type NetStats struct {
	BytesSent uint64 `json:"bytes_sent"`
	BytesRecv uint64 `json:"bytes_recv"`
}

func main() {
	// --- 1. Parse command line arguments ---
	pidFlag := flag.Int("pid", 0, "The process ID (PID) to monitor")
	intervalFlag := flag.Int("interval", 5, "Monitoring interval in seconds")
	flag.Parse()

	if *pidFlag <= 0 {
		log.Fatal("Please provide a valid PID with the -pid flag.")
	}

	pid := *pidFlag
	interval := time.Duration(*intervalFlag) * time.Second

	log.Printf("Monitoring process with PID: %d, interval: %v\n", pid, interval)

	// --- 2. Initial network stats for calculating delta ---
	prevStats, err := getSystemNetStats()
	if err != nil {
		log.Fatalf("Failed to get initial network stats: %v", err)
	}

	// --- 3. Main monitoring loop ---
	for {
		// Clear the screen (for better readability)
		// Note: This works on Unix-like systems. For Windows, use "cls".
		fmt.Print("\033[H\033[2J")

		// --- A. Get and print process connections ---
		connections, err := getProcessConnections(pid)
		if err != nil {
			log.Printf("Error getting connections: %v", err)
		} else {
			fmt.Println("--- Network Connections ---")
			if len(connections) == 0 {
				fmt.Println("No active network connections found.")
			} else {
				// Print as JSON for clarity
				jsonConnections, _ := json.MarshalIndent(connections, "", "  ")
				fmt.Println(string(jsonConnections))
			}
		}
		fmt.Println() // Add a blank line

		// --- B. Get and print listening ports ---
		listeningPorts, err := getListeningPorts(pid)
		if err != nil {
			log.Printf("Error getting listening ports: %v", err)
		} else {
			fmt.Println("--- Listening Ports ---")
			if len(listeningPorts) == 0 {
				fmt.Println("Not listening on any ports.")
			} else {
				for _, port := range listeningPorts {
					fmt.Printf("Port: %d\n", port)
				}
			}
		}
		fmt.Println() // Add a blank line

		// --- C. Get and print network traffic (system-wide) ---
		currentStats, err := getSystemNetStats()
		if err != nil {
			log.Printf("Error getting current network stats: %v", err)
		} else {
			// Calculate the delta (change) since last check
			deltaSent := currentStats.BytesSent - prevStats.BytesSent
			deltaRecv := currentStats.BytesRecv - prevStats.BytesRecv

			fmt.Println("--- Network Traffic (System Wide) ---")
			fmt.Printf("Sent:     %s/s\n", formatBytes(deltaSent))
			fmt.Printf("Received: %s/s\n", formatBytes(deltaRecv))

			// Update previous stats for next iteration
			prevStats = currentStats
		}

		time.Sleep(interval)
	}
}

// getProcessConnections retrieves all network connections for a given PID.
// It works by parsing the /proc/[pid]/net/tcp and /proc/[pid]/net/udp files.
func getProcessConnections(pid int) ([]Connection, error) {
	var connections []Connection

	// Check if the process is still alive
	procPath := fmt.Sprintf("/proc/%d", pid)
	if _, err := os.Stat(procPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("process with PID %d does not exist", pid)
	}

	// We need to check both TCP and UDP connections
	protocols := []struct {
		name string
		file string
	}{
		{"tcp", "net/tcp"},
		{"udp", "net/udp"},
	}

	for _, proto := range protocols {
		filePath := filepath.Join(procPath, proto.file)
		file, err := os.Open(filePath)
		if err != nil {
			// If the file doesn't exist, the process might not have that type of connection
			if os.IsNotExist(err) {
				continue
			}
			return nil, fmt.Errorf("failed to open %s: %w", filePath, err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		// Skip the header line (sl, local_address, rem_address, ...)
		scanner.Scan()

		for scanner.Scan() {
			line := scanner.Text()
			parts := strings.Fields(line)
			if len(parts) < 10 {
				continue // Malformed line, skip
			}

			localAddr := parts[1]
			remoteAddr := parts[2]
			status := parts[3]

			localIP, localPort := parseIPPort(localAddr)
			remoteIP, remotePort := parseIPPort(remoteAddr)

			connections = append(connections, Connection{
				Protocol:   strings.ToUpper(proto.name),
				LocalIP:    localIP,
				LocalPort:  localPort,
				RemoteIP:   remoteIP,
				RemotePort: remotePort,
				Status:     getTcpStatus(status),
			})
		}

		if err := scanner.Err(); err != nil {
			log.Printf("Error scanning %s: %v", filePath, err)
		}
	}

	return connections, nil
}

// parseIPPort converts a string like "0100007F:0050" to IP and port.
// The IP part is in hex and in reverse byte order.
func parseIPPort(addr string) (string, int) {
	ipPort := strings.Split(addr, ":")
	if len(ipPort) != 2 {
		return "N/A", 0
	}

	// Parse IP (reverse the bytes)
	ipHex := ipPort[0]
	var ipParts []string
	for i := 0; i < 8; i += 2 {
		part := ipHex[i : i+2]
		ipParts = append([]string{part}, ipParts...) // Prepend to reverse
	}
	ip, _ := strconv.ParseInt(strings.Join(ipParts, ""), 16, 64)

	// Parse Port
	port, _ := strconv.ParseInt(ipPort[1], 16, 64)

	return fmt.Sprintf("%d.%d.%d.%d",
		byte(ip>>24), byte(ip>>16), byte(ip>>8), byte(ip)), int(port)
}

// getTcpStatus converts a hex status code to a human-readable string.
func getTcpStatus(hexCode string) string {
	// For a complete list, see: https://github.com/torvalds/linux/blob/master/include/net/tcp_states.h
	statusMap := map[string]string{
		"01": "ESTABLISHED",
		"02": "SYN_SENT",
		"03": "SYN_RECV",
		"04": "FIN_WAIT1",
		"05": "FIN_WAIT2",
		"06": "TIME_WAIT",
		"07": "CLOSE",
		"08": "CLOSE_WAIT",
		"09": "LAST_ACK",
		"0A": "LISTEN",
		"0B": "CLOSING",
	}
	if status, ok := statusMap[hexCode]; ok {
		return status
	}
	return "UNKNOWN"
}

// getListeningPorts extracts all unique local ports that the process is listening on.
func getListeningPorts(pid int) ([]int, error) {
	connections, err := getProcessConnections(pid)
	if err != nil {
		return nil, err
	}

	listeningPorts := make(map[int]bool)
	for _, conn := range connections {
		if conn.Status == "LISTEN" {
			listeningPorts[conn.LocalPort] = true
		}
	}

	// Convert map keys to slice
	ports := make([]int, 0, len(listeningPorts))
	for port := range listeningPorts {
		ports = append(ports, port)
	}

	return ports, nil
}

// getSystemNetStats uses gopsutil to get total bytes sent/received.
func getSystemNetStats() (NetStats, error) {
	ioCounters, err := net.IOCounters(false) // false for all interfaces combined
	if err != nil {
		return NetStats{}, err
	}
	if len(ioCounters) == 0 {
		return NetStats{}, fmt.Errorf("no network interfaces found")
	}
	counter := ioCounters[0]
	return NetStats{
		BytesSent: counter.BytesSent,
		BytesRecv: counter.BytesRecv,
	}, nil
}

// formatBytes converts bytes to a human-readable string (e.g., KB, MB).
func formatBytes(bytes uint64) string {
	const unit = 1024
	if bytes < unit {
		return fmt.Sprintf("%d B", bytes)
	}
	div, exp := uint64(unit), 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.2f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
}
